import pandas as pd
import os

output_path = './output'

file_list = sorted([i for i in os.listdir(output_path) if '最终结果' in i])
print(file_list)

file_name = os.path.join(output_path, file_list[-1])

df = pd.read_excel(file_name, dtype={'股票代码': str})

print(df.head())
print(df.columns)


# 处理境外债券列，去除包含无明细|未披露|无信息|无金额的数据
print("处理前数据行数:", len(df))

# 创建过滤条件
filter_keywords = ['无明细', '未披露', '无信息', '无金额', '金额：无', '金额：-', '金额：0.00', '金额：0元', '金额：0.0', '金额：不适用']
# 同时处理QDII列和境外债券列
# 创建两个过滤条件
mask_qdii = ~df['QDII'].astype(str).str.contains('|'.join(filter_keywords), na=False)
mask_foreign_bond = ~df['境外债券'].astype(str).str.contains('|'.join(filter_keywords), na=False)

# 组合条件：两列不能同时包含关键词
mask = mask_qdii | mask_foreign_bond
print(f"QDII列过滤前数据行数: {len(df[mask_qdii])}")
print(f"境外债券列过滤前数据行数: {len(df[mask_foreign_bond])}")
print(f"两列组合过滤前数据行数: {len(df[mask])}")

# 应用过滤条件
df_filtered = df[mask]

print("处理后数据行数:", len(df_filtered))
print("已移除包含以下关键词的记录:", filter_keywords)

# 保存处理后的数据
# 读取原始文件中的基本信息sheet
df_basic_info = pd.read_excel(file_name, sheet_name='基本信息', dtype={'证券代码': str})

# 获取过滤后的股票代码列表
filtered_codes = df_filtered['股票代码'].unique().tolist()

# 从基本信息中筛选出包含在过滤后数据中的证券代码的记录
filtered_basic_info = df_basic_info[df_basic_info['证券代码'].apply(lambda x: x.split('.')[0] in filtered_codes)]

# 创建过滤后的文件名
filtered_file_name = os.path.join(output_path, '过滤后_' + os.path.basename(file_name))

# 使用ExcelWriter保存多个sheet
with pd.ExcelWriter(filtered_file_name) as writer:
    # 保存解析结果sheet
    df_filtered.to_excel(writer, sheet_name='解析结果', index=False)
    # 保存基本信息sheet
    filtered_basic_info.to_excel(writer, sheet_name='基本信息', index=False)

print(f"已将过滤后的数据保存至: {filtered_file_name}")
print(f"解析结果sheet: {len(df_filtered)}条记录")
print(f"基本信息sheet: {len(filtered_basic_info)}条记录")

